{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1fcc525-8149-451e-b9f6-f1656a406d39",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install burr[start]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f7d4516-c5e1-4889-abed-155645c41d9d",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Traces/Spans\n",
    "\n",
    "This walks through the traces/spans example. You can see the code in [application.py](application.py).\n",
    "\n",
    "To ground this, you can see an example action implementation (from `application.py`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "82986ee3-3662-4a74-b4d1-cced1caa1095",
   "metadata": {},
   "outputs": [],
   "source": [
    "from burr.core import action, State\n",
    "from burr.visibility import TracerFactory\n",
    "\n",
    "from typing import Tuple\n",
    "\n",
    "@action(reads=[\"prompt\"], writes=[\"mode\"])\n",
    "def choose_mode(state: State, __tracer: TracerFactory) -> Tuple[dict, State]:\n",
    "    with __tracer(\"generate_prompt\"):\n",
    "        prompt = (\n",
    "            f\"You are a chatbot. You've been prompted this: {state['prompt']}. \"\n",
    "            f\"You have the capability of responding in the following modes: {', '.join(MODES)}. \"\n",
    "            \"Please respond with *only* a single word representing the mode that most accurately \"\n",
    "            \"corresponds to the prompt. Fr instance, if the prompt is 'draw a picture of a cat', \"\n",
    "            \"the mode would be 'generate_image'. If the prompt is 'what is the capital of France', the mode would be 'answer_question'.\"\n",
    "            \"If none of these modes apply, please respond with 'unknown'.\"\n",
    "        )\n",
    "    with __tracer(\"query_openai\", span_dependencies=[\"generate_prompt\"]):\n",
    "        with __tracer(\"create_openai_client\"):\n",
    "            client = _get_openai_client()\n",
    "        with __tracer(\"query_openai\"):\n",
    "            result = client.chat.completions.create(\n",
    "                model=\"gpt-4\",\n",
    "                messages=[\n",
    "                    {\"role\": \"system\", \"content\": \"You are a helpful assistant\"},\n",
    "                    {\"role\": \"user\", \"content\": prompt},\n",
    "                ],\n",
    "            )\n",
    "    with __tracer(\"process_openai_response\", span_dependencies=[\"query_openai\"]):\n",
    "        content = result.choices[0].message.content\n",
    "        mode = content.lower()\n",
    "        if mode not in MODES:\n",
    "            mode = \"unknown\"\n",
    "        result = {\"mode\": mode}\n",
    "    return result, state.update(**result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a07bdcf3-94f4-47d0-8ec3-00d3c7ee8c8d",
   "metadata": {},
   "source": [
    "# Running the code\n",
    "\n",
    "The following runs the application on a simple prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bba29bd1-3e59-4050-821c-8da44c2c583f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from application import application as tracing_and_spans_application"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "db048ea6-dd1a-4088-93e3-508cb5d86c9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "app = tracing_and_spans_application()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a7c9586a-28a8-4f84-a2a2-6028a699371e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 8.1.0 (20230707.0739)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"586pt\" height=\"410pt\"\n",
       " viewBox=\"0.00 0.00 586.22 410.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 406)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-406 582.22,-406 582.22,4 -4,4\"/>\n",
       "<!-- prompt -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M532,-335C532,-335 500.25,-335 500.25,-335 494.25,-335 488.25,-329 488.25,-323 488.25,-323 488.25,-311 488.25,-311 488.25,-305 494.25,-299 500.25,-299 500.25,-299 532,-299 532,-299 538,-299 544,-305 544,-311 544,-311 544,-323 544,-323 544,-329 538,-335 532,-335\"/>\n",
       "<text text-anchor=\"middle\" x=\"516.12\" y=\"-311.95\" font-family=\"Times,serif\" font-size=\"14.00\">prompt</text>\n",
       "</g>\n",
       "<!-- check_safety -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>check_safety</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M512.38,-268C512.38,-268 449.88,-268 449.88,-268 443.88,-268 437.88,-262 437.88,-256 437.88,-256 437.88,-244 437.88,-244 437.88,-238 443.88,-232 449.88,-232 449.88,-232 512.38,-232 512.38,-232 518.38,-232 524.38,-238 524.38,-244 524.38,-244 524.38,-256 524.38,-256 524.38,-262 518.38,-268 512.38,-268\"/>\n",
       "<text text-anchor=\"middle\" x=\"481.12\" y=\"-244.95\" font-family=\"Times,serif\" font-size=\"14.00\">check_safety</text>\n",
       "</g>\n",
       "<!-- prompt&#45;&gt;check_safety -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>prompt&#45;&gt;check_safety</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M506.75,-298.58C503.29,-292.17 499.3,-284.75 495.51,-277.71\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"498.35,-276.61 490.53,-269.47 492.19,-279.93 498.35,-276.61\"/>\n",
       "</g>\n",
       "<!-- input__prompt -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>input__prompt</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" cx=\"516.12\" cy=\"-384\" rx=\"62.1\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"516.12\" y=\"-378.95\" font-family=\"Times,serif\" font-size=\"14.00\">input: prompt</text>\n",
       "</g>\n",
       "<!-- input__prompt&#45;&gt;prompt -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>input__prompt&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M516.12,-365.58C516.12,-359.61 516.12,-352.76 516.12,-346.15\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"519.63,-346.47 516.12,-336.47 512.63,-346.47 519.63,-346.47\"/>\n",
       "</g>\n",
       "<!-- decide_mode -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M355.5,-185.5C355.5,-185.5 290.75,-185.5 290.75,-185.5 284.75,-185.5 278.75,-179.5 278.75,-173.5 278.75,-173.5 278.75,-161.5 278.75,-161.5 278.75,-155.5 284.75,-149.5 290.75,-149.5 290.75,-149.5 355.5,-149.5 355.5,-149.5 361.5,-149.5 367.5,-155.5 367.5,-161.5 367.5,-161.5 367.5,-173.5 367.5,-173.5 367.5,-179.5 361.5,-185.5 355.5,-185.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"323.12\" y=\"-162.45\" font-family=\"Times,serif\" font-size=\"14.00\">decide_mode</text>\n",
       "</g>\n",
       "<!-- check_safety&#45;&gt;decide_mode -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>check_safety&#45;&gt;decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M446.84,-231.53C423.45,-219.61 392.25,-203.72 367.03,-190.87\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"368.84,-187.35 358.34,-185.93 365.66,-193.59 368.84,-187.35\"/>\n",
       "<text text-anchor=\"middle\" x=\"444.88\" y=\"-203.7\" font-family=\"Times,serif\" font-size=\"14.00\">safe=True</text>\n",
       "</g>\n",
       "<!-- response -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M400.75,-36C400.75,-36 361.5,-36 361.5,-36 355.5,-36 349.5,-30 349.5,-24 349.5,-24 349.5,-12 349.5,-12 349.5,-6 355.5,0 361.5,0 361.5,0 400.75,0 400.75,0 406.75,0 412.75,-6 412.75,-12 412.75,-12 412.75,-24 412.75,-24 412.75,-30 406.75,-36 400.75,-36\"/>\n",
       "<text text-anchor=\"middle\" x=\"381.12\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">response</text>\n",
       "</g>\n",
       "<!-- check_safety&#45;&gt;response -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>check_safety&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M488.95,-231.69C503.29,-197.13 529.31,-118.48 495.12,-67 479.27,-43.13 448.98,-31.1 423.6,-25.06\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"424.59,-21.48 414.08,-22.82 423.13,-28.32 424.59,-21.48\"/>\n",
       "</g>\n",
       "<!-- generate_image -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M90.25,-103C90.25,-103 12,-103 12,-103 6,-103 0,-97 0,-91 0,-91 0,-79 0,-79 0,-73 6,-67 12,-67 12,-67 90.25,-67 90.25,-67 96.25,-67 102.25,-73 102.25,-79 102.25,-79 102.25,-91 102.25,-91 102.25,-97 96.25,-103 90.25,-103\"/>\n",
       "<text text-anchor=\"middle\" x=\"51.12\" y=\"-79.95\" font-family=\"Times,serif\" font-size=\"14.00\">generate_image</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_image -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M278.27,-164.27C205.67,-160.22 68.95,-150.53 53.88,-134.5 48.81,-129.11 46.89,-121.73 46.56,-114.35\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"50.1,-114.56 47.07,-104.4 43.11,-114.23 50.1,-114.56\"/>\n",
       "<text text-anchor=\"middle\" x=\"116.75\" y=\"-121.2\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_image</text>\n",
       "</g>\n",
       "<!-- generate_code -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M218.12,-103C218.12,-103 148.12,-103 148.12,-103 142.12,-103 136.12,-97 136.12,-91 136.12,-91 136.12,-79 136.12,-79 136.12,-73 142.12,-67 148.12,-67 148.12,-67 218.12,-67 218.12,-67 224.12,-67 230.12,-73 230.12,-79 230.12,-79 230.12,-91 230.12,-91 230.12,-97 224.12,-103 218.12,-103\"/>\n",
       "<text text-anchor=\"middle\" x=\"183.12\" y=\"-79.95\" font-family=\"Times,serif\" font-size=\"14.00\">generate_code</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_code -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M278.4,-161.16C244.57,-156.16 201.96,-147.49 190.12,-134.5 185.1,-128.99 182.7,-121.56 181.71,-114.18\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"185.16,-114.07 181.14,-104.27 178.17,-114.44 185.16,-114.07\"/>\n",
       "<text text-anchor=\"middle\" x=\"248.62\" y=\"-121.2\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_code</text>\n",
       "</g>\n",
       "<!-- answer_question -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M342.5,-103C342.5,-103 259.75,-103 259.75,-103 253.75,-103 247.75,-97 247.75,-91 247.75,-91 247.75,-79 247.75,-79 247.75,-73 253.75,-67 259.75,-67 259.75,-67 342.5,-67 342.5,-67 348.5,-67 354.5,-73 354.5,-79 354.5,-79 354.5,-91 354.5,-91 354.5,-97 348.5,-103 342.5,-103\"/>\n",
       "<text text-anchor=\"middle\" x=\"301.12\" y=\"-79.95\" font-family=\"Times,serif\" font-size=\"14.00\">answer_question</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;answer_question -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M318.35,-149.03C315.55,-138.8 311.96,-125.64 308.77,-113.99\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"311.9,-113.16 305.89,-104.43 305.15,-115 311.9,-113.16\"/>\n",
       "<text text-anchor=\"middle\" x=\"379\" y=\"-121.2\" font-family=\"Times,serif\" font-size=\"14.00\">mode=answer_question</text>\n",
       "</g>\n",
       "<!-- prompt_for_more -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M473.88,-103C473.88,-103 384.38,-103 384.38,-103 378.38,-103 372.38,-97 372.38,-91 372.38,-91 372.38,-79 372.38,-79 372.38,-73 378.38,-67 384.38,-67 384.38,-67 473.88,-67 473.88,-67 479.88,-67 485.88,-73 485.88,-79 485.88,-79 485.88,-91 485.88,-91 485.88,-97 479.88,-103 473.88,-103\"/>\n",
       "<text text-anchor=\"middle\" x=\"429.12\" y=\"-79.95\" font-family=\"Times,serif\" font-size=\"14.00\">prompt_for_more</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;prompt_for_more -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M367.96,-160.85C399.65,-155.83 438.2,-147.29 448.12,-134.5 453.06,-128.14 452.27,-120.34 449.13,-112.87\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"451.82,-111.49 443.95,-104.39 445.68,-114.85 451.82,-111.49\"/>\n",
       "</g>\n",
       "<!-- generate_image&#45;&gt;response -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>generate_image&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M102.59,-72.25C110.78,-70.44 119.18,-68.64 127.12,-67 201.29,-51.7 288.04,-35.72 338.56,-26.6\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"339.16,-29.87 348.38,-24.65 337.92,-22.98 339.16,-29.87\"/>\n",
       "</g>\n",
       "<!-- generate_code&#45;&gt;response -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>generate_code&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M230.55,-68.43C263.55,-57.6 307.14,-43.29 338.93,-32.85\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"339.82,-35.92 348.23,-29.47 337.64,-29.27 339.82,-35.92\"/>\n",
       "</g>\n",
       "<!-- answer_question&#45;&gt;response -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>answer_question&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M322.56,-66.58C331.35,-59.44 341.68,-51.05 351.19,-43.32\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"353.07,-45.49 358.63,-36.47 348.66,-40.06 353.07,-45.49\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more&#45;&gt;response -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>prompt_for_more&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M416.26,-66.58C411.42,-60.02 405.8,-52.41 400.49,-45.23\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"402.78,-43.43 394.02,-37.47 397.15,-47.59 402.78,-43.43\"/>\n",
       "</g>\n",
       "<!-- response&#45;&gt;prompt -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>response&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M412.89,-24.37C462.34,-33.65 552.12,-54.54 552.12,-84 552.12,-251 552.12,-251 552.12,-251 552.12,-264.91 545.74,-278.63 538.35,-289.86\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"535.02,-287.49 532.02,-297.65 540.7,-291.58 535.02,-287.49\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x11ace2d90>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "app.visualize(include_conditions=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "58131e17-6f48-4d8b-b6c6-32605c30d62a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'chat_item': {'content': 'Silent stars above,\\n'\n",
      "                          'Whispers of the gentle breeze,\\n'\n",
      "                          'Peaceful dreams unfold.',\n",
      "               'role': 'assistant',\n",
      "               'type': 'text'}}\n"
     ]
    }
   ],
   "source": [
    "import pprint\n",
    "\n",
    "action, result, state = app.run(halt_after=[\"response\"], inputs={\"prompt\" : \"please write a haiku\"})\n",
    "pprint.pprint(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c47d87b3-8d32-420b-ab75-acf2a7f6d2a9",
   "metadata": {},
   "source": [
    "# Viewing in UI\n",
    "\n",
    "The following cell gives you a link to the prior run in the UI. If you have not yet, run `burr` in your terminal to start the local tracking server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "93cad8f7-9f0b-4b20-bc75-2adaca5bfde9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "[Link to UI](http://localhost:7241/project/demo_tracing/d50acb12-080d-499c-b1d5-719862b8ae58)"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Markdown\n",
    "url = f\"[Link to UI](http://localhost:7241/project/demo_tracing/{app.uid})\"\n",
    "Markdown(url)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
